package myUtils

import (
	"KD_web/go/dao"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"net/http"
	"strconv"
	"time"
)

func RateLimiter(limit int, window time.Duration) gin.HandlerFunc {
	// 根据ip进行限流
	return func(c *gin.Context) {
		userIP := c.ClientIP()
		key := "rate_limit:" + userIP
		// 获取当前ip请求次数
		count, err := dao.RedisClient.Get(dao.Ctx, key).Result()
		if err != nil && err != redis.Nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Internal Server Error"})
			c.Abort()
			return
		}
		currentCount, _ := strconv.Atoi(count)
		if currentCount >= limit {
			c.JSON(http.StatusTooManyRequests, gin.H{"error": "Too Many Requests"})
			c.Abort()
			return
		}
		// 增加计数
		if err := dao.RedisClient.Incr(dao.Ctx, key).Err(); err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Internal Server Error"})
			c.Abort()
			return
		}
		// 设置过期时间
		if currentCount == 0 {
			if err := dao.RedisClient.Expire(dao.Ctx, key, window).Err(); err != nil {
				c.JSON(http.StatusInternalServerError, gin.H{"error": "Internal Server Error"})
				c.Abort()
				return
			}
		}

		c.Next()
	}
}
